#define _POSIX_C_SOURCE 200112L
#define _DARWIN_C_SOURCE
#include "memory_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <errno.h>
#include <sys/mman.h>
#include <inttypes.h>

// Helper functions
static uint64_t get_current_time_us(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000 + ts.tv_nsec / 1000;
}

static void* aligned_malloc(size_t size, size_t alignment) {
    void* ptr;
    if (posix_memalign(&ptr, alignment, size) != 0) {
        return NULL;
    }
    return ptr;
}

// Memory manager initialization
memory_manager_t* mm_init(size_t total_memory) {
    memory_manager_t* mm = calloc(1, sizeof(memory_manager_t));
    if (!mm) {
        return NULL;
    }

    // Initialize global mutex
    if (pthread_mutex_init(&mm->global_mutex, NULL) != 0) {
        free(mm);
        return NULL;
    }

    // Initialize memory pools for different types
    const char* pool_names[] = {
        "PageBuffer", "WALBuffer", "QueryBuffer", 
        "IndexBuffer", "TempBuffer"
    };
    
    size_t pool_sizes[] = {
        total_memory * 0.6,  // 60% for page buffers
        total_memory * 0.2,  // 20% for WAL buffers
        total_memory * 0.1,  // 10% for query buffers
        total_memory * 0.05, // 5% for index buffers
        total_memory * 0.05  // 5% for temp buffers
    };

    for (int i = 0; i < MEM_TYPE_COUNT && i < MEM_POOL_MAX_POOLS; i++) {
        if (pool_init(&mm->pools[i], pool_names[i], pool_sizes[i]) != 0) {
            // Cleanup on failure
            for (int j = 0; j < i; j++) {
                pool_cleanup(&mm->pools[j]);
            }
            pthread_mutex_destroy(&mm->global_mutex);
            free(mm);
            return NULL;
        }
    }

    // Initialize buffer pool
    mm->buffer_pool = calloc(1, sizeof(buffer_pool_t));
    if (!mm->buffer_pool) {
        for (int i = 0; i < MEM_TYPE_COUNT; i++) {
            pool_cleanup(&mm->pools[i]);
        }
        pthread_mutex_destroy(&mm->global_mutex);
        free(mm);
        return NULL;
    }

    if (bp_init(mm->buffer_pool, BUFFER_POOL_MAX_PAGES) != 0) {
        free(mm->buffer_pool);
        for (int i = 0; i < MEM_TYPE_COUNT; i++) {
            pool_cleanup(&mm->pools[i]);
        }
        pthread_mutex_destroy(&mm->global_mutex);
        free(mm);
        return NULL;
    }

    // Start background flush thread
    mm->flush_thread_running = true;
    if (pthread_create(&mm->flush_thread, NULL, flush_thread_func, mm) != 0) {
        mm->flush_thread_running = false;
        bp_cleanup(mm->buffer_pool);
        free(mm->buffer_pool);
        for (int i = 0; i < MEM_TYPE_COUNT; i++) {
            pool_cleanup(&mm->pools[i]);
        }
        pthread_mutex_destroy(&mm->global_mutex);
        free(mm);
        return NULL;
    }

    return mm;
}

// Memory manager cleanup
int mm_cleanup(memory_manager_t* mm) {
    if (!mm) {
        return 0;
    }

    // Stop flush thread
    mm->flush_thread_running = false;
    if (mm->flush_thread) {
        pthread_join(mm->flush_thread, NULL);
    }

    // Cleanup buffer pool
    if (mm->buffer_pool) {
        bp_cleanup(mm->buffer_pool);
        free(mm->buffer_pool);
    }

    // Cleanup memory pools
    for (int i = 0; i < MEM_POOL_MAX_POOLS; i++) {
        pool_cleanup(&mm->pools[i]);
    }

    pthread_mutex_destroy(&mm->global_mutex);
    free(mm);
    return 0;
}

// Memory allocation
void* mm_alloc(memory_manager_t* mm, size_t size, mem_type_t type) {
    if (!mm || size == 0 || type >= MEM_TYPE_COUNT) {
        return NULL;
    }

    pthread_mutex_lock(&mm->global_mutex);
    mm->alloc_requests++;

    void* ptr = pool_alloc(&mm->pools[type], size);
    if (ptr) {
        mm->total_allocated += size;
        if (mm->total_allocated > mm->peak_usage) {
            mm->peak_usage = mm->total_allocated;
        }
    } else {
        mm->oom_failures++;
        // Try to handle memory pressure
        pthread_mutex_unlock(&mm->global_mutex);
        if (mm_handle_pressure(mm) == 0) {
            pthread_mutex_lock(&mm->global_mutex);
            ptr = pool_alloc(&mm->pools[type], size);
            if (ptr) {
                mm->total_allocated += size;
                if (mm->total_allocated > mm->peak_usage) {
                    mm->peak_usage = mm->total_allocated;
                }
            }
        } else {
            pthread_mutex_lock(&mm->global_mutex);
        }
    }

    pthread_mutex_unlock(&mm->global_mutex);
    return ptr;
}

// Memory free
void mm_free(memory_manager_t* mm, void* ptr) {
    if (!mm || !ptr) {
        return;
    }

    pthread_mutex_lock(&mm->global_mutex);
    mm->free_requests++;

    // Find which pool this pointer belongs to
    for (int i = 0; i < MEM_POOL_MAX_POOLS; i++) {
        mem_block_t* block = mm->pools[i].used_list;
        while (block) {
            if (block->ptr == ptr) {
                mm->total_allocated -= block->size;
                pool_free(&mm->pools[i], ptr);
                pthread_mutex_unlock(&mm->global_mutex);
                return;
            }
            block = block->next;
        }
    }

    pthread_mutex_unlock(&mm->global_mutex);
}

// Aligned allocation
void* mm_aligned_alloc(memory_manager_t* mm, size_t size, size_t alignment, mem_type_t type) {
    if (!mm || size == 0 || type >= MEM_TYPE_COUNT) {
        return NULL;
    }

    // For now, use regular allocation with alignment padding
    size_t padded_size = size + alignment - 1;
    void* ptr = mm_alloc(mm, padded_size, type);
    if (!ptr) {
        return NULL;
    }

    // Align the pointer
    uintptr_t addr = (uintptr_t)ptr;
    uintptr_t aligned_addr = (addr + alignment - 1) & ~(alignment - 1);
    
    return (void*)aligned_addr;
}

// Buffer pool initialization
int bp_init(buffer_pool_t* bp, uint32_t num_pages) {
    if (!bp || num_pages == 0) {
        return -1;
    }

    memset(bp, 0, sizeof(buffer_pool_t));

    // Initialize mutex and condition variable
    if (pthread_mutex_init(&bp->mutex, NULL) != 0) {
        return -1;
    }
    
    if (pthread_cond_init(&bp->flush_cond, NULL) != 0) {
        pthread_mutex_destroy(&bp->mutex);
        return -1;
    }

    // Allocate page array
    bp->pages = calloc(num_pages, sizeof(buffer_page_t));
    if (!bp->pages) {
        pthread_cond_destroy(&bp->flush_cond);
        pthread_mutex_destroy(&bp->mutex);
        return -1;
    }

    // Allocate hash table
    bp->hash_table = calloc(BUFFER_POOL_HASH_SIZE, sizeof(buffer_page_t*));
    if (!bp->hash_table) {
        free(bp->pages);
        pthread_cond_destroy(&bp->flush_cond);
        pthread_mutex_destroy(&bp->mutex);
        return -1;
    }

    bp->total_pages = num_pages;

    // Initialize page data
    for (uint32_t i = 0; i < num_pages; i++) {
        bp->pages[i].page_data = aligned_malloc(GUDB_PAGE_SIZE, GUDB_PAGE_SIZE);
        if (!bp->pages[i].page_data) {
            // Cleanup on failure
            for (uint32_t j = 0; j < i; j++) {
                free(bp->pages[j].page_data);
            }
            free(bp->pages);
            free(bp->hash_table);
            pthread_cond_destroy(&bp->flush_cond);
            pthread_mutex_destroy(&bp->mutex);
            return -1;
        }
        bp->pages[i].page_id = UINT32_MAX; // Invalid page ID
    }

    return 0;
}

// Buffer pool cleanup
int bp_cleanup(buffer_pool_t* bp) {
    if (!bp) {
        return 0;
    }

    // Flush all dirty pages
    bp_flush_all_dirty(bp);

    // Free page data
    if (bp->pages) {
        for (uint32_t i = 0; i < bp->total_pages; i++) {
            if (bp->pages[i].page_data) {
                free(bp->pages[i].page_data);
            }
        }
        free(bp->pages);
    }

    if (bp->hash_table) {
        free(bp->hash_table);
    }

    pthread_cond_destroy(&bp->flush_cond);
    pthread_mutex_destroy(&bp->mutex);

    return 0;
}

// Get page from buffer pool
buffer_page_t* bp_get_page(buffer_pool_t* bp, uint32_t page_id) {
    if (!bp) {
        return NULL;
    }

    pthread_mutex_lock(&bp->mutex);

    // Check if page is already in buffer pool
    buffer_page_t* page = hash_lookup(bp, page_id);
    if (page) {
        // Cache hit
        bp->cache_hits++;
        page->access_count++;
        page->last_access_time = get_current_time_us();
        lru_promote(bp, page);
        pthread_mutex_unlock(&bp->mutex);
        return page;
    }

    // Cache miss
    bp->cache_misses++;

    // Find a free page or evict one
    page = NULL;
    for (uint32_t i = 0; i < bp->total_pages; i++) {
        if (bp->pages[i].page_id == UINT32_MAX) {
            page = &bp->pages[i];
            break;
        }
    }

    if (!page) {
        // Need to evict a page
        page = lru_evict_candidate(bp);
        if (!page) {
            pthread_mutex_unlock(&bp->mutex);
            return NULL;
        }

        // Remove from hash table and LRU list
        hash_remove(bp, page);
        lru_remove(bp, page);

        if (page->dirty) {
            // Flush dirty page (simplified - in real implementation, 
            // this would write to disk)
            bp->writes++;
            page->dirty = false;
            bp->dirty_pages--;
        }

        bp->evictions++;
    }

    // Initialize the page
    page->page_id = page_id;
    page->dirty = false;
    page->pinned = false;
    page->pin_count = 0;
    page->access_count = 1;
    page->last_access_time = get_current_time_us();

    // Add to hash table and LRU list
    hash_insert(bp, page);
    lru_insert_head(bp, page);

    bp->used_pages++;
    bp->reads++;

    pthread_mutex_unlock(&bp->mutex);
    return page;
}

// Put page back to buffer pool
int bp_put_page(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    // Update access information
    page->last_access_time = get_current_time_us();
    lru_promote(bp, page);

    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// Pin a page in memory
int bp_pin_page(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    if (!page->pinned) {
        page->pinned = true;
        bp->pinned_pages++;
    }
    page->pin_count++;
    
    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// Unpin a page
int bp_unpin_page(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    if (page->pinned && page->pin_count > 0) {
        page->pin_count--;
        if (page->pin_count == 0) {
            page->pinned = false;
            bp->pinned_pages--;
        }
    }
    
    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// Mark page as dirty
int bp_mark_dirty(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    if (!page->dirty) {
        page->dirty = true;
        bp->dirty_pages++;
        
        // Signal flush thread if too many dirty pages
        if (bp->dirty_pages >= BUFFER_POOL_MAX_DIRTY) {
            pthread_cond_signal(&bp->flush_cond);
        }
    }
    
    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// Flush a specific page
int bp_flush_page(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    if (page->dirty) {
        // In real implementation, this would write to disk
        bp->writes++;
        page->dirty = false;
        bp->dirty_pages--;
    }
    
    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// Flush all dirty pages
int bp_flush_all_dirty(buffer_pool_t* bp) {
    if (!bp) {
        return -1;
    }

    pthread_mutex_lock(&bp->mutex);
    
    for (uint32_t i = 0; i < bp->total_pages; i++) {
        if (bp->pages[i].dirty) {
            // In real implementation, this would write to disk
            bp->writes++;
            bp->pages[i].dirty = false;
        }
    }
    
    bp->dirty_pages = 0;
    
    pthread_mutex_unlock(&bp->mutex);
    return 0;
}

// LRU operations
int lru_promote(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    // Remove from current position
    lru_remove(bp, page);
    
    // Insert at head
    lru_insert_head(bp, page);
    
    return 0;
}

buffer_page_t* lru_evict_candidate(buffer_pool_t* bp) {
    if (!bp) {
        return NULL;
    }

    // Find the least recently used unpinned page
    buffer_page_t* page = bp->lru_tail;
    while (page && page->pinned) {
        page = page->lru_prev;
    }
    
    return page;
}

int lru_insert_head(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    page->lru_next = bp->lru_head;
    page->lru_prev = NULL;
    
    if (bp->lru_head) {
        bp->lru_head->lru_prev = page;
    } else {
        bp->lru_tail = page;
    }
    
    bp->lru_head = page;
    return 0;
}

int lru_remove(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    if (page->lru_prev) {
        page->lru_prev->lru_next = page->lru_next;
    } else {
        bp->lru_head = page->lru_next;
    }
    
    if (page->lru_next) {
        page->lru_next->lru_prev = page->lru_prev;
    } else {
        bp->lru_tail = page->lru_prev;
    }
    
    page->lru_next = page->lru_prev = NULL;
    return 0;
}

// Hash table operations
uint32_t hash_page_id(uint32_t page_id) {
    // Simple hash function
    return page_id % BUFFER_POOL_HASH_SIZE;
}

buffer_page_t* hash_lookup(buffer_pool_t* bp, uint32_t page_id) {
    if (!bp) {
        return NULL;
    }

    uint32_t hash = hash_page_id(page_id);
    buffer_page_t* page = bp->hash_table[hash];
    
    while (page) {
        if (page->page_id == page_id) {
            return page;
        }
        page = page->hash_next;
    }
    
    return NULL;
}

int hash_insert(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    uint32_t hash = hash_page_id(page->page_id);
    
    page->hash_next = bp->hash_table[hash];
    page->hash_prev = NULL;
    
    if (bp->hash_table[hash]) {
        bp->hash_table[hash]->hash_prev = page;
    }
    
    bp->hash_table[hash] = page;
    return 0;
}

int hash_remove(buffer_pool_t* bp, buffer_page_t* page) {
    if (!bp || !page) {
        return -1;
    }

    uint32_t hash = hash_page_id(page->page_id);
    
    if (page->hash_prev) {
        page->hash_prev->hash_next = page->hash_next;
    } else {
        bp->hash_table[hash] = page->hash_next;
    }
    
    if (page->hash_next) {
        page->hash_next->hash_prev = page->hash_prev;
    }
    
    page->hash_next = page->hash_prev = NULL;
    return 0;
}

// Memory pool operations
int pool_init(mem_pool_t* pool, const char* name, size_t size) {
    if (!pool || !name || size == 0) {
        return -1;
    }

    memset(pool, 0, sizeof(mem_pool_t));
    strncpy(pool->name, name, sizeof(pool->name) - 1);
    pool->total_size = size;

    if (pthread_mutex_init(&pool->mutex, NULL) != 0) {
        return -1;
    }

    return 0;
}

int pool_cleanup(mem_pool_t* pool) {
    if (!pool) {
        return 0;
    }

    // Free all blocks
    mem_block_t* block = pool->used_list;
    while (block) {
        mem_block_t* next = block->next;
        free(block->ptr);
        free(block);
        block = next;
    }

    block = pool->free_list;
    while (block) {
        mem_block_t* next = block->next;
        free(block->ptr);
        free(block);
        block = next;
    }

    pthread_mutex_destroy(&pool->mutex);
    return 0;
}

void* pool_alloc(mem_pool_t* pool, size_t size) {
    if (!pool || size == 0) {
        return NULL;
    }

    pthread_mutex_lock(&pool->mutex);

    // Align size
    size = (size + MEM_POOL_ALIGNMENT - 1) & ~(MEM_POOL_ALIGNMENT - 1);

    // Check if we have enough space
    if (pool->used_size + size > pool->total_size) {
        pthread_mutex_unlock(&pool->mutex);
        return NULL;
    }

    // Allocate new block
    mem_block_t* block = malloc(sizeof(mem_block_t));
    if (!block) {
        pthread_mutex_unlock(&pool->mutex);
        return NULL;
    }

    block->ptr = malloc(size);
    if (!block->ptr) {
        free(block);
        pthread_mutex_unlock(&pool->mutex);
        return NULL;
    }

    block->size = size;
    block->in_use = true;
    block->next = pool->used_list;
    block->prev = NULL;

    if (pool->used_list) {
        pool->used_list->prev = block;
    }
    pool->used_list = block;

    pool->used_size += size;
    pool->alloc_count++;

    if (pool->used_size > pool->peak_usage) {
        pool->peak_usage = pool->used_size;
    }

    pthread_mutex_unlock(&pool->mutex);
    return block->ptr;
}

void pool_free(mem_pool_t* pool, void* ptr) {
    if (!pool || !ptr) {
        return;
    }

    pthread_mutex_lock(&pool->mutex);

    // Find the block
    mem_block_t* block = pool->used_list;
    while (block) {
        if (block->ptr == ptr) {
            // Remove from used list
            if (block->prev) {
                block->prev->next = block->next;
            } else {
                pool->used_list = block->next;
            }
            
            if (block->next) {
                block->next->prev = block->prev;
            }

            pool->used_size -= block->size;
            pool->free_count++;

            free(block->ptr);
            free(block);
            break;
        }
        block = block->next;
    }

    pthread_mutex_unlock(&pool->mutex);
}

size_t pool_available(const mem_pool_t* pool) {
    if (!pool) {
        return 0;
    }
    return pool->total_size - pool->used_size;
}

// Background flush thread
void* flush_thread_func(void* arg) {
    memory_manager_t* mm = (memory_manager_t*)arg;
    
    while (mm->flush_thread_running) {
        sleep(1); // Flush every second
        
        if (mm->buffer_pool && mm->buffer_pool->dirty_pages > 0) {
            // Flush some dirty pages
            pthread_mutex_lock(&mm->buffer_pool->mutex);
            
            uint32_t pages_to_flush = mm->buffer_pool->dirty_pages / 4; // Flush 25%
            if (pages_to_flush == 0) pages_to_flush = 1;
            
            uint32_t flushed = 0;
            for (uint32_t i = 0; i < mm->buffer_pool->total_pages && flushed < pages_to_flush; i++) {
                if (mm->buffer_pool->pages[i].dirty && !mm->buffer_pool->pages[i].pinned) {
                    mm->buffer_pool->pages[i].dirty = false;
                    mm->buffer_pool->dirty_pages--;
                    mm->buffer_pool->writes++;
                    flushed++;
                }
            }
            
            pthread_mutex_unlock(&mm->buffer_pool->mutex);
        }
    }
    
    return NULL;
}

// Handle memory pressure
int mm_handle_pressure(memory_manager_t* mm) {
    if (!mm) {
        return -1;
    }

    // Try to free unused pages from buffer pool
    if (mm->buffer_pool) {
        pthread_mutex_lock(&mm->buffer_pool->mutex);
        
        // Flush all dirty pages
        bp_flush_all_dirty(mm->buffer_pool);
        
        // Try to evict some clean pages
        uint32_t evicted = 0;
        for (uint32_t i = 0; i < mm->buffer_pool->total_pages && evicted < 10; i++) {
            if (!mm->buffer_pool->pages[i].pinned && !mm->buffer_pool->pages[i].dirty) {
                if (mm->buffer_pool->pages[i].page_id != UINT32_MAX) {
                    hash_remove(mm->buffer_pool, &mm->buffer_pool->pages[i]);
                    lru_remove(mm->buffer_pool, &mm->buffer_pool->pages[i]);
                    mm->buffer_pool->pages[i].page_id = UINT32_MAX;
                    mm->buffer_pool->used_pages--;
                    evicted++;
                }
            }
        }
        
        pthread_mutex_unlock(&mm->buffer_pool->mutex);
        return evicted > 0 ? 0 : -1;
    }
    
    return -1;
}

// Statistics
void mm_print_stats(const memory_manager_t* mm) {
    if (!mm) {
        return;
    }

    printf("Memory Manager Statistics:\n");
    printf("  Total allocated: %zu bytes\n", mm->total_allocated);
    printf("  Peak usage: %zu bytes\n", mm->peak_usage);
    printf("  Allocation requests: %" PRIu64 "\n", mm->alloc_requests);
    printf("  Free requests: %" PRIu64 "\n", mm->free_requests);
    printf("  OOM failures: %" PRIu64 "\n", mm->oom_failures);

    printf("\nMemory Pools:\n");
    for (int i = 0; i < MEM_POOL_MAX_POOLS; i++) {
        const mem_pool_t* pool = &mm->pools[i];
        if (pool->total_size > 0) {
            printf("  %s: %zu/%zu bytes (%.1f%% used)\n", 
                   pool->name, pool->used_size, pool->total_size,
                   (double)pool->used_size / pool->total_size * 100.0);
        }
    }

    if (mm->buffer_pool) {
        bp_print_stats(mm->buffer_pool);
    }
}

void bp_print_stats(const buffer_pool_t* bp) {
    if (!bp) {
        return;
    }

    printf("\nBuffer Pool Statistics:\n");
    printf("  Total pages: %u\n", bp->total_pages);
    printf("  Used pages: %u\n", bp->used_pages);
    printf("  Dirty pages: %u\n", bp->dirty_pages);
    printf("  Pinned pages: %u\n", bp->pinned_pages);
    printf("  Cache hits: %" PRIu64 "\n", bp->cache_hits);
    printf("  Cache misses: %" PRIu64 "\n", bp->cache_misses);
    printf("  Hit ratio: %.2f%%\n", bp_get_hit_ratio(bp) * 100.0);
    printf("  Reads: %" PRIu64 "\n", bp->reads);
    printf("  Writes: %" PRIu64 "\n", bp->writes);
    printf("  Evictions: %" PRIu64 "\n", bp->evictions);
}

double bp_get_hit_ratio(const buffer_pool_t* bp) {
    if (!bp || (bp->cache_hits + bp->cache_misses) == 0) {
        return 0.0;
    }
    return (double)bp->cache_hits / (bp->cache_hits + bp->cache_misses);
}

size_t mm_get_usage(const memory_manager_t* mm, mem_type_t type) {
    if (!mm || type >= MEM_TYPE_COUNT) {
        return 0;
    }
    return mm->pools[type].used_size;
}